# files: data/feature_loader.py
import pandas as pd
from datetime import timedelta
from data.stock_data_loader import StockDataLoader
from data.feature_generator import FeatureGenerator
from data.hit_aggregator import aggregate_by_stock_date

class FeatureLoader:
    def __init__(self, data_loader=None):
        self.loader = data_loader or StockDataLoader()
        self.deal_dates = self.loader.get_deal_dates()["add_date"].tolist()

    def get_next_trade_date(self, current_date):
        for i, date in enumerate(self.deal_dates):
            if date == current_date and i + 1 < len(self.deal_dates):
                return self.deal_dates[i + 1]
        return None

    def get_trade_date_range(self, start_date, end_date, extend_days=60):
        """从交易日表中获取范围内的交易日，用于日线数据提取"""
        if start_date not in self.deal_dates or end_date not in self.deal_dates:
            raise ValueError("开始或结束日期不在交易日列表中")

        start_idx = self.deal_dates.index(start_date)
        end_idx = self.deal_dates.index(end_date)

        from_idx = max(0, start_idx - extend_days)
        to_idx = min(len(self.deal_dates) - 1, end_idx + 5)
        return self.deal_dates[from_idx], self.deal_dates[to_idx]

    def merge_hit_with_next_day_increase(self, start_date=None, end_date=None):
        # 加载 hits
        hits_df = self.loader.get_hit_board_data(start_date, end_date)
        if hits_df.empty:
            return pd.DataFrame()

        # next_trade_date 标签
        hits_df["next_trade_date"] = hits_df["add_date"].apply(self.get_next_trade_date)

        # 加载日线数据（仅命中范围）
        date_from, date_to = self.get_trade_date_range(start_date, end_date, extend_days=60)
        daily_df = self.loader.get_daily_summary_data(start_date=date_from, end_date=date_to)

        # 合并 next_day 标签
        merged = pd.merge(
            hits_df,
            daily_df,
            left_on=["stock_code", "next_trade_date"],
            right_on=["stock_code", "summary_date"],
            suffixes=("", "_next")
        )

        merged["next_day_high_premium"] = merged["increase_next"] > 5.0
        return merged

    def get_training_data(self, start_date=None, end_date=None):
        merged_df = self.merge_hit_with_next_day_increase(start_date, end_date)

        # ⏬ 聚合 hits 成个股-日期级别
        agg_df = aggregate_by_stock_date(merged_df)

        # ⏬ 合并聚合结果回原始 merged 数据
        final_df = pd.merge(
            merged_df,
            agg_df,
            on=["stock_code", "add_date"],
            how="left"
        )

        # ⏬ 衍生特征
        final_df = FeatureGenerator(final_df).generate()

        return final_df
